/**
 * @file main.cpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-10-26
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#include <dlfcn.h>
#include <lyn_api.h>
#include <nlohmann/json.hpp>
#include <string>
#include "imageDecoder.h"
#include "imageEncoder.h"
#include "ipeParamYolov8.h"
#include "opencvWindow.hpp"
#include "osdPlugin.h"
#include "segment_post_process.h"
#include "util.hpp"
#include "CLI/CLI11.hpp"

using namespace std;

using SEGMENT_POST_FUNC = int (*)(SegmentPostProcessInfo_t *);

int main(int argc, char *argv[])
{
  CLI::App app;

  string inputFile;
  string outputFile;
  int deviceId;
  string postPlugin;
  string yolov8Model;
  string osdPlugin;

  app.add_option("input_image", inputFile, "input jpeg file path")
      ->default_val("../data/yolov8_segment_with_image_sample/test.jpeg");
  app.add_option("output_image", outputFile, "output jpeg file path")
      ->default_val("../output.jpeg");
  app.add_option("device_id", deviceId, "lynxi device id")
      ->default_val(0);
  app.add_option("-m,--yolov8_model", yolov8Model, "yolov8 model path")
      ->default_val("../model/yolov8s-seg_b1_h640_w640_c3_iuint8_ofloat32_0000/Net_0");
  app.add_option("-p,--post_plugin", postPlugin, "yolov8 post process plugin path")
      ->default_val("../common/plugin/postprocess/lib/libYolov8SegmentPostPlugin.so");
  app.add_option("-o,--osd_plugin", osdPlugin, "osd process plugin path")
      ->default_val("../common/plugin/osd-opencv/lib/libOsdPlugin.so");

  CLI11_PARSE(app, argc, argv);

  if (!IsFileExist(inputFile))
  {
    cout << "error : input jpeg file is not exsit." << endl;
    return 0;
  }

  if (!IsValidDevice(deviceId))
  {
    cout << "error : device id is invalid." << endl;
    return 0;
  }

  if (!IsFileExist(postPlugin))
  {
    cout << "error : yolov8 post process plugin is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(yolov8Model))
  {
    cout << "error : yolov8 model is not exsit." << endl;
    return 0;
  }

  if (!IsFileExist(osdPlugin))
  {
    cout << "error : osd plugin is not exsit." << endl;
    return 0;
  }

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, deviceId));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  // 3. 初始化解码图片类
  ImageDecoder imageDec;
  imageDec.Init(inputFile);
  lynImageInfo_t imgInfo;
  imageDec.GetImageDecInfo(&imgInfo);

  // 4. 同步解码图片到 Device 侧
  lynCodecBuf_t decodeImg;
  decodeImg.size = imgInfo.output.predictBufSize;
  CHECK_ERR(lynMalloc((void **)&decodeImg.data, decodeImg.size));
  imageDec.DecodeImageToDevice(stream, decodeImg.data);

  // 5. 加载模型
  ModelInfo yolov8ModelInfo;
  yolov8ModelInfo.LoadModelByPath(yolov8Model.c_str());
  lynPixelFormat_t pixelFormat = LYN_PIX_FMT_NV12;

  // 6. IPE 处理
  uint8_t *pIpeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pIpeBufOut, yolov8ModelInfo.inputSize));
  IpeParamYolov8 ipeYolov8(yolov8ModelInfo.width, yolov8ModelInfo.height);
  ipeYolov8.SetImgInfo(imgInfo.output.width, imgInfo.output.height,
                       pixelFormat);
  ipeYolov8.CalcParam(stream, decodeImg.data, pIpeBufOut);

  // 7. 调用 APU 推理接口
  uint8_t *pDevApuOutData = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevApuOutData, yolov8ModelInfo.outputSize));
  CHECK_ERR(lynExecuteModelAsync(stream, yolov8ModelInfo.model, pIpeBufOut,
                                 pDevApuOutData, yolov8ModelInfo.batchSize));

  // 8. 加载post Plugin, 进行模型后处理
  SegmentPostProcessInfo_t post_info;
  memset(&post_info, 0, sizeof(SegmentPostProcessInfo_t));

  post_info.model_width = yolov8ModelInfo.width;
  post_info.model_height = yolov8ModelInfo.height;
  post_info.img_width = imgInfo.output.width;
  post_info.img_height = imgInfo.output.height;
  post_info.padx = ipeYolov8.m_iPadX;
  post_info.pady = ipeYolov8.m_iPadY;
  post_info.fatio = ipeYolov8.m_fRatio;

  lynSegBoxesInfo *pDevBoxesInfo;
  CHECK_ERR(lynMalloc((void **)&pDevBoxesInfo, sizeof(lynSegBoxesInfo)));
  uint8_t *devMaskData;

  // 后处理
  {
    void *handle = dlopen(postPlugin.c_str(), RTLD_NOW | RTLD_GLOBAL);
    if (handle == nullptr)
    {
      cout << "error : post plugin load fail. error:" << dlerror() << endl;
      return 0;
    }
    SEGMENT_POST_FUNC segmentPostFunc =
        (SEGMENT_POST_FUNC)dlsym(handle, "lynSegmentPostProcess");
    if (segmentPostFunc == nullptr)
    {
      cout << "error : post plugin has no lynSegmentPostProcess function."
           << endl;
      dlclose(handle);
      return 0;
    }

    lynSynchronizeStream(stream);
    uint8_t *pHostApuOutData = new uint8_t[yolov8ModelInfo.outputSize];
    CHECK_ERR(lynMemcpy(pHostApuOutData, pDevApuOutData,
                        yolov8ModelInfo.outputSize, ServerToClient));

    post_info.output_tensor = pHostApuOutData;
    post_info.boxesInfo = new lynSegBoxesInfo();
    post_info.boxesInfo->maskDataLen = 0;
    post_info.boxesInfo->maskData = nullptr;

    segmentPostFunc(&post_info);
    CHECK_ERR(lynMemcpy(pDevBoxesInfo, post_info.boxesInfo, sizeof(lynSegBoxesInfo),
                        ClientToServer));

    if (post_info.boxesInfo->maskDataLen > 0)
    {
      CHECK_ERR(lynMalloc((void **)&devMaskData, post_info.boxesInfo->maskDataLen));
      CHECK_ERR(lynMemcpy(devMaskData, post_info.boxesInfo->maskData, post_info.boxesInfo->maskDataLen, ClientToServer));
      delete[] post_info.boxesInfo->maskData;
    }

    delete post_info.boxesInfo;
    delete[] pHostApuOutData;
    dlclose(handle);
  }

  // 9. 加载osd Plugin
  {
    // 画框
    lynPlugin_t plugin;
    CHECK_ERR(lynPluginRegister(&plugin, osdPlugin.c_str()));

    lynDrawSegBoxAndTextPara para;
    para.imgData = decodeImg.data;
    para.imgFmt = LYN_PIX_FMT_NV12;
    para.imgW = imgInfo.output.width;
    para.imgH = imgInfo.output.height;
    para.boxesInfo = pDevBoxesInfo;
    para.boxColor = DRAW_COLOR_BLUE;
    para.boxThick = DRAW_THICK_2;
    para.fontSize = FONT_SIZE_24;
    para.fontColor = DRAW_COLOR_BLUE;
    CHECK_ERR(lynPluginRunAsync(stream, plugin, "lynDrawSegBoxAndText", &para,
                                sizeof(para)));
    lynSynchronizeStream(stream);

    // 画分隔图
    // nv12tobgr
    uint8_t *pImgBgr = nullptr;
    int imgBgrSize = imgInfo.output.width * imgInfo.output.height * 3;
    CHECK_ERR(lynMalloc((void **)&pImgBgr, imgBgrSize));
    IpeParamC2C ipeBgr(LYN_PIX_FMT_BGR24);
    ipeBgr.SetImgInfo(imgInfo.output.width, imgInfo.output.height, pixelFormat);
    ipeBgr.CalcParam(stream, decodeImg.data, pImgBgr);

    MASK_PARAM_T osdParam;
    osdParam.imgData = pImgBgr;
    osdParam.width = imgInfo.output.width;
    osdParam.height = imgInfo.output.height;
    osdParam.boxInfo = pDevBoxesInfo;
    osdParam.maskData = devMaskData;

    CHECK_ERR(lynPluginRunAsync(stream, plugin, "maskDraw", &osdParam,
                                sizeof(osdParam)));
    // bgrtonv12
    IpeParamC2C ipeNv12(LYN_PIX_FMT_NV12);
    ipeNv12.SetImgInfo(imgInfo.output.width, imgInfo.output.height, LYN_PIX_FMT_BGR24);
    ipeNv12.CalcParam(stream, pImgBgr, decodeImg.data);

    lynSynchronizeStream(stream);

    CHECK_ERR(lynFree(pImgBgr));
    CHECK_ERR(lynFree(pDevBoxesInfo));
    CHECK_ERR(lynPluginUnregister(plugin));
  }

  // 等待 plugin 处理完，释放内存
  CHECK_ERR(lynFree(pDevApuOutData));
  CHECK_ERR(lynFree(pIpeBufOut));

  // 9. 进行图片编码
  ImageEncoder encoder;
  encoder.EncodeImageAndSave(stream, outputFile, decodeImg, imgInfo);

  // 10. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynFree(decodeImg.data));
  yolov8ModelInfo.UnLoadModel();

  if (ctx)
  {
    CHECK_ERR(lynDestroyContext(ctx));
  }

  // delete[] labelList.labels;

  return 0;
}
